#include "Base64Tools.h"

std::string Base64Tools::base64Encode(const unsigned char *message, const size_t length) {
  if (message == NULL || length <= 0)
    return "";
  int encodedSize = 4 * ceil((double)length / 3);
  std::vector<char> encodedMessage(encodedSize + 1);

  BIO *b64 = BIO_new(BIO_f_base64());
  BIO *bio = BIO_new(BIO_s_mem());
  bio = BIO_push(b64, bio);
  BIO_set_flags(bio, BIO_FLAGS_BASE64_NO_NL);

  if (BIO_write(bio, message, length) <= 0) {
    BIO_free_all(bio);
    return "";
  }
  if (BIO_flush(bio) <= 0) {
    BIO_free_all(bio);
    return "";
  }

  BUF_MEM *bufferPtr;
  BIO_get_mem_ptr(bio, &bufferPtr);
  if (BIO_set_close(bio, BIO_NOCLOSE) <= 0) {
    std::cerr << "Warning: BIO_set_close failed" << std::endl;
  }

  // Add a NUL terminator
  memcpy(encodedMessage.data(), bufferPtr->data, bufferPtr->length);
  encodedMessage[bufferPtr->length] = '\0';

  BIO_free_all(bio);
  return std::string(encodedMessage.data());;
}

std::vector<unsigned char> Base64Tools::base64Decode(const char *encodedMessage, const size_t encodedMessageLength) {
  if (encodedMessage == NULL || encodedMessageLength <= 0)
    return std::vector<unsigned char>();
  int decodedLength = calculateDecodedLength(encodedMessage, encodedMessageLength);
  std::vector<unsigned char> decodedMessage(decodedLength + 1);

  BIO *bio = BIO_new_mem_buf(encodedMessage, encodedMessageLength);
  BIO *b64 = BIO_new(BIO_f_base64());
  bio = BIO_push(b64, bio);
  BIO_set_flags(bio, BIO_FLAGS_BASE64_NO_NL);

  // Add a NUL terminator
  decodedLength = BIO_read(bio, decodedMessage.data(), encodedMessageLength);
  if (decodedLength < 0) {
    BIO_free_all(bio);
    return std::vector<unsigned char>();
  }
  decodedMessage[decodedLength] = '\0';

  BIO_free_all(bio);
  decodedMessage.resize(decodedLength);
  return decodedMessage;
}

int Base64Tools::calculateDecodedLength(const char *encodedMessage, const size_t encodedMessageLength) {
  if (encodedMessage == NULL || encodedMessageLength <= 0)
    return 0;
  unsigned int padding = 0;

  // Check for trailing '=''s as padding
  if(encodedMessage[encodedMessageLength - 1] == '=' && encodedMessage[encodedMessageLength - 2] == '=') {
    padding = 2;
  } else if (encodedMessage[encodedMessageLength - 1] == '=') {
    padding = 1;
  }

  return static_cast<int>(encodedMessageLength) * 0.75 - padding;
}
